#include "kernel_operator.h"

using namespace AscendC;

constexpr int32_t BUFFER_NUM = 2;

template<typename T>
class KernelAddBf16 {
public:
    __aicore__ inline KernelAddBf16() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR z,
                                uint32_t totalSizeX,
                                uint32_t totalSizeY,
                                uint32_t totalTileNum,
                                uint32_t tileLength)
    {
        this->totalSizeX = totalSizeX;
        this->totalSizeY = totalSizeY;
        this->totalTileNum = totalTileNum;
        this->tileLength = tileLength;
        this->typeSize = sizeof(T);

        xGm.GSetBuffer((__gm__ T *)x, totalSizeX);
        yGm.GSetBuffer((__gm__ T *)y, totalSizeY);
        zGm.GSetBuffer((__gm__ T *)z, totalSizeX);
        pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileLength * this->typeSize);
        pipe.InitBuffer(inQueueY, BUFFER_NUM, this->tileLength * this->typeSize);
        pipe.InitBuffer(outQueueZ, BUFFER_NUM, this->tileLength * this->typeSize);

        pipe.InitBuffer(xTempBuffer, tileLength * sizeof(float));
        pipe.InitBuffer(yTempBuffer, tileLength * sizeof(float));
        pipe.InitBuffer(zTempBuffer, tileLength * sizeof(float));
    }
    __aicore__ inline void Process()
    {
        uint32_t block_id = GetBlockIdx();
        uint32_t core_num = GetBlockNum();
        uint32_t repNum = this->totalSizeX / this->totalSizeY;

        for (uint32_t rep = 0; rep < repNum; rep++) {
            this->repOffset = rep * this->totalSizeY;
            this->copyBytes = this->tileLength * this->typeSize;
            for (uint32_t i = block_id; i < this->totalTileNum; i += core_num) {
                if (i == this->totalTileNum - 1) { // last tile need pad
                    uint32_t lastTileLength = this->totalSizeY - (this->totalTileNum - 1) * this->tileLength;
                    this->copyBytes = lastTileLength * this->typeSize;
                }

                CopyIn(i);
                Compute();
                CopyOut(i);
            }
        }
    }

private:
    __aicore__ inline void CopyIn(int32_t progress)
    {
        LocalTensor<T> xLocal = inQueueX.AllocTensor<T>();
        LocalTensor<T> yLocal = inQueueY.AllocTensor<T>();
        DataCopyExtParams  copyParams = {1, this->copyBytes, 0, 0, 0};
        DataCopyPadExtParams<T> copyPadParams = {true, 0, 0, 0};
        DataCopyPad(xLocal, xGm[this->repOffset + progress * this->tileLength], copyParams, copyPadParams);
        DataCopyPad(yLocal, yGm[progress * this->tileLength], copyParams, copyPadParams);
        inQueueX.EnQue(xLocal);
        inQueueY.EnQue(yLocal);
    }
    __aicore__ inline void Compute()
    {
        LocalTensor<T> xLocal = inQueueX.DeQue<T>();
        LocalTensor<T> yLocal = inQueueY.DeQue<T>();
        LocalTensor<T> zLocal = outQueueZ.AllocTensor<T>();

        LocalTensor<float> xTempLocal = xTempBuffer.Get<float>();
        LocalTensor<float> yTempLocal = yTempBuffer.Get<float>();
        LocalTensor<float> zTempLocal = zTempBuffer.Get<float>();
        // bf16 -> f32
        Cast(xTempLocal, xLocal, RoundMode::CAST_NONE, this->tileLength);
        Cast(yTempLocal, yLocal, RoundMode::CAST_NONE, this->tileLength);
        pipe_barrier(PIPE_V);
        inQueueX.FreeTensor(xLocal);
        inQueueY.FreeTensor(yLocal);

        Add(zTempLocal, xTempLocal, yTempLocal, this->tileLength);
        pipe_barrier(PIPE_V);

        // f32 -> bf16
        Cast(zLocal, zTempLocal, RoundMode::CAST_RINT, this->tileLength);
        pipe_barrier(PIPE_V);
        outQueueZ.EnQue<T>(zLocal);
    }
    __aicore__ inline void CopyOut(int32_t progress)
    {
        LocalTensor<T> zLocal = outQueueZ.DeQue<T>();
        DataCopyExtParams copyParams = {1, this->copyBytes, 0, 0, 0};
        DataCopyPad(zGm[this->repOffset + progress * this->tileLength], zLocal, copyParams);
        outQueueZ.FreeTensor(zLocal);
    }

private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX, inQueueY;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueZ;

    TBuf<TPosition::VECCALC> xTempBuffer;
    TBuf<TPosition::VECCALC> yTempBuffer;
    TBuf<TPosition::VECCALC> zTempBuffer;
    GlobalTensor<T> xGm;
    GlobalTensor<T> yGm;
    GlobalTensor<T> zGm;

    uint32_t totalSizeX;
    uint32_t totalSizeY;
    uint32_t totalTileNum;
    uint32_t tileLength;
    uint32_t copyBytes;
    uint32_t repOffset{0};
    uint32_t typeSize;
};

